-- 令牌桶算法的Lua脚本实现
-- 确保在分布式环境下的原子性操作
-- 
-- KEYS[1]: 令牌桶的key前缀
-- ARGV[1]: 最大令牌数
-- ARGV[2]: 令牌生成速率（每秒）
-- ARGV[3]: 当前时间戳（毫秒）
-- ARGV[4]: 请求令牌数量（默认为1）

-- 获取参数
local key = KEYS[1]
local tokensKey = key .. ":tokens"
local timestampKey = key .. ":timestamp"

local maxTokens = tonumber(ARGV[1])
local tokensPerSecond = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requestTokens = tonumber(ARGV[4]) or 1

-- 获取当前令牌数和上次更新时间
local tokensStr = redis.call('get', tokensKey)
local timestampStr = redis.call('get', timestampKey)

-- 如果key不存在，初始化为满额令牌
local currentTokens = maxTokens
local lastRefillTime = now

if tokensStr then
    currentTokens = tonumber(tokensStr)
    lastRefillTime = tonumber(timestampStr)
end

-- 计算从上次更新到现在应该新增的令牌数
local timePassed = math.max(0, now - lastRefillTime)
local newTokens = timePassed / 1000.0 * tokensPerSecond

-- 更新当前令牌数，不超过最大令牌数
currentTokens = math.min(maxTokens, currentTokens + newTokens)

-- 如果令牌数足够，则获取成功
if currentTokens >= requestTokens then
    -- 消耗令牌
    currentTokens = currentTokens - requestTokens
    
    -- 更新令牌数和时间戳
    redis.call('set', tokensKey, tostring(currentTokens))
    redis.call('set', timestampKey, tostring(now))
    
    -- 设置过期时间，避免长期占用Redis内存
    -- 计算令牌桶重新填满需要的时间的两倍作为过期时间
    local expireTime = math.ceil(maxTokens / tokensPerSecond) * 2
    redis.call('expire', tokensKey, expireTime)
    redis.call('expire', timestampKey, expireTime)
    
    return 1 -- 获取令牌成功
else
    return 0 -- 获取令牌失败，令牌不足
end